package cn.mrdjun.redis.stater.ratelimiter;

import com.google.common.collect.ImmutableList;
import lombok.RequiredArgsConstructor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.lang.reflect.Method;
import java.util.Objects;

/**
 * 令牌桶限流算法
 * 拦截器检测是否存在注解，有注解则对当前接口限流控制
 *
 * @author DengJun 2021/3/23
 */
@RequiredArgsConstructor
public class RateLimiterInterceptor implements HandlerInterceptor {
    private final static Logger log = LoggerFactory.getLogger(RateLimiterInterceptor.class);
    private final RedisTemplate<String, Object> redisTemplate;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (handler instanceof HandlerMethod) {
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            Method method = handlerMethod.getMethod();
            RateLimit rateLimit = AnnotationUtils.findAnnotation(method, RateLimit.class);
            // 尝试获取方法、类上的限流注解
            if (Objects.isNull(rateLimit)) {
                rateLimit = AnnotationUtils.findAnnotation(handlerMethod.getBean().getClass(), RateLimit.class);
            }

            // 若无注解 - 直接放行
            if (Objects.isNull(rateLimit)) {
                return true;
            }

            // 尝试获取令牌 - 若未获取到令牌
            if (!tryAcquire(request.getRequestURI(), rateLimit.capacity(), rateLimit.rate())) {
                log.warn("超出频率限制：接口[{}]令牌桶容量[{}]速率[{}]", request.getRequestURI(), rateLimit.capacity(), rateLimit.rate());
                return false;
            }
        }
        return true;
    }

    /**
     * lua脚本，限流
     */
    private final static String TEXT = "local ratelimit_info = redis.pcall('HMGET',KEYS[1],'last_time','current_token')\n" +
            "local last_time = ratelimit_info[1]\n" +
            "local current_token = tonumber(ratelimit_info[2])\n" +
            "local max_token = tonumber(ARGV[1])\n" +
            "local token_rate = tonumber(ARGV[2])\n" +
            "local current_time = tonumber(ARGV[3])\n" +
            "if current_token == nil then\n" +
            "  current_token = max_token\n" +
            "  last_time = current_time\n" +
            "else\n" +
            "  local past_time = current_time-last_time\n" +
            "  \n" +
            "  if past_time>1000 then\n" +
            "\t  current_token = current_token+token_rate\n" +
            "\t  last_time = current_time\n" +
            "  end\n" +
            "\n" +
            "  if current_token>max_token then\n" +
            "    current_token = max_token\n" +
            "\tlast_time = current_time\n" +
            "  end\n" +
            "end\n" +
            "\n" +
            "local result = 0\n" +
            "if(current_token>0) then\n" +
            "  result = 1\n" +
            "  current_token = current_token-1\n" +
            "  last_time = current_time\n" +
            "end\n" +
            "redis.call('HMSET',KEYS[1],'last_time',last_time,'current_token',current_token)\n" +
            "return result";


    /**
     * 获取令牌
     *
     * @param key  请求ID或者请求路径
     * @param max  最大能同时承受多少的并发（桶容量）
     * @param rate 每秒生成多少的令牌
     * @return 获取令牌返回true，没有获取返回false
     */
    public boolean tryAcquire(String key, int max, int rate) {
        ImmutableList<String> keys = ImmutableList.of(key);
        RedisScript<Long> script = new DefaultRedisScript<>(TEXT, Long.class);
        return Long.valueOf(1).equals(redisTemplate.execute(script, keys, max, rate, System.currentTimeMillis()));
    }
}
